from celery import Celery, Task
from kombu import serialization
from datetime import datetime

from app.celery_task.models.celery_task import CeleryTask
from common.database import SessionLocal
from core.logger import logger, celery_task_id
from setting import config

celery_app = Celery(
    'worker',
    backend=config.CELERY_RESULT_BACKEND,  # 存储结果
    broker=config.CELERY_BROKER_URL,  # 消息中间件
)

serialization.registry._decoders.pop("application/x-python-serialize")

celery_app.conf.update(
    accept_content=['json'],
    task_serializer='json',
    result_serializer='json',
    enable_utc=False,
    timezone='Asia/Shanghai'
)


class TaskMonitor(Task):

    def before_start(self, task_id, args, kwargs):
        celery_task_id.set(str(task_id))
        session = SessionLocal()

        # 记录数据库
        type = self.name.split(".")[1]
        task_name = self.name.split(".")[-1]
        comment = f"{task_name} begin"
        if args and isinstance(args[0], dict):
            relation_uuid = args[0].get("uuid", 0)
        else:
            relation_uuid = 0
        task_data = {
            "name": task_name,
            "type": type,
            "tid": task_id,
            "relation_uuid": relation_uuid,
            "state": "STARTED",
            "progress": 0,
            "start_at": datetime.now(),
            "comment": comment
        }
        result = CeleryTask.add(session, task_data)
        if not result:
            logger.error(f"before_start=>task_name:{task_name} task_id:{task_id}, add task failed")
            return False
        
        session.close()
        logger.info(f"before_start=>task_name:{task_name} task_id:{task_id}, add task success")
        return True

    def on_failure(self, exc, task_id, args, kwargs, einfo):
        """
        failed callback
        任务执行失败，有异常抛出
        """
        logger.info(f"on_failure=> exc:{exc}, task_id:{task_id}, args:{args}, kwargs:{kwargs}, einfo:{einfo}")
        # session = SessionLocal()
        # task_name = self.name.split(".")[-1]
        # comment = f"{task_name} failed"
        #
        # logger.info(f"on_failure=>task_name:{task_name} task_id:{task_id}, einfo:{einfo}")
        #
        # result = CeleryTask.update_task(session, task_id, state="FAILURE", comment=comment)
        # if not result:
        #     logger.info(f"on_failure=>task_name:{task_name} task_id:{task_id}, update task status:FAILURE failed")
        #     return False
        
        # session.close()
        # logger.info(f"on_failure=>task_name:{task_name} task_id:{task_id}, update task status:FAILURE success")
        # return False

    def on_success(self, retval, task_id, args, kwargs):
        """
        success callback
        任务函数跑完，有返回值，需要判断返回值为True 或者False
        """
        logger.info(f"on_success=> retval:{retval}, task_id:{task_id}, args:{args}, kwargs:{kwargs}")
        # session = SessionLocal()
        #
        # task_name = self.name.split(".")[-1]
        # if retval:
        #     comment = f"{task_name} success"
        #     state = "SUCCESS"
        # else:
        #     comment = f"{task_name} failed"
        #     state = "FAILURE"
        # result = CeleryTask.update_task(session, task_id, state=state, comment=comment)
        # if not result:
        #     logger.info(f"on_success=>task_name:{task_name} task_id:{task_id}, update task status:{state} failed")
        #     return False
        #
        # session.close()
        # logger.info(f"on_success=>task_name:{task_name} task_id:{task_id}, update task status:{state} success")
        # return True

    def on_retry(self, exc, task_id, args, kwargs, einfo):
        """retry callback"""
        logger.info(f"on_retry=> exc:{exc}, task_id:{task_id}, args:{args}, kwargs:{kwargs}, einfo:{einfo}")

    def after_return(self, status, retval, task_id, args, kwargs, einfo):
        """Handler called after the task returns."""
        # logger.info(f"on_retry=> status:{status}, retval:{retval}, task_id:{task_id}, args:{args}, kwargs:{kwargs}, einfo:{einfo}")
        session = SessionLocal()

        task_name = self.name.split(".")[-1]
        if status == "SUCCESS" and retval == None:
            # 任务函数正常跑完，但是没有return
            comment = f"{task_name} success"
            state = "SUCCESS"
        elif status == "FAILURE" or not retval:
            # 函数抛了异常，或者返回值为假
            comment = f"{task_name} failed"
            state = "FAILURE"
        else:
            comment = f"{task_name} success"
            state = "SUCCESS"

        result = CeleryTask.update_task(session, task_id, state=state, comment=comment)
        if not result:
            logger.info(f"after_return=>task_name:{task_name} task_id:{task_id}, update task status:{state} failed")
            return False
        
        session.close()
        logger.info(f"after_return=>task_name:{task_name} task_id:{task_id}, update task status:{state} success")
        return True
